#include <jni.h>
#include <string>
#include <cstdio>
#include <ctime>

#ifdef __cplusplus
extern "C" {
#endif
#include "../ffmpeg/include/libavcodec/avcodec.h"
#include "../ffmpeg/include/libavformat/avformat.h"
#include "../ffmpeg/include/libavutil/log.h"
#include "../ffmpeg/include/libavutil/time.h"
#ifdef __cplusplus
} // endof extern "C"
#endif

#ifdef ANDROID

#include <android/log.h>

#define LOGE(format, ...)  __android_log_print(ANDROID_LOG_ERROR, "(>_<)", format, ##__VA_ARGS__)
#define LOGI(format, ...)  __android_log_print(ANDROID_LOG_INFO,  "(^_^)", format, ##__VA_ARGS__)
#else
#define LOGE(format, ...)  printf("(>_<) " format "\n", ##__VA_ARGS__)
#define LOGI(format, ...)  printf("(^_^) " format "\n", ##__VA_ARGS__)
#endif

//Output FFmpeg's av_log()
void custom_log(void *ptr, int level, const char *fmt, va_list vl) {

    //To TXT file
    FILE *fp = fopen("/storage/emulated/0/av_log.txt", "a+");
    if (fp) {
        vfprintf(fp, fmt, vl);
        fflush(fp);
        fclose(fp);
    }
}

extern "C"
JNIEXPORT jint JNICALL
Java_com_example_sffmpegandroidstreamer_MainActivity_pushStream(JNIEnv *env, jobject thiz,
                                                                jstring input_url,
                                                                jstring output_url,
                                                                jobject callback) {

    jobject resultObj = env->NewLocalRef(callback);
    jclass cls = env->GetObjectClass(resultObj);
    jmethodID jmethodId_Result = env->GetMethodID(cls, "onResult", "(ILjava/lang/String;)V");

    AVOutputFormat *ofmt = nullptr;
    AVFormatContext *ifmt_ctx = nullptr, *ofmt_ctx = nullptr;
    AVPacket pkt;

    int ret, i;
    char input_str[500] = {0};
    char output_str[500] = {0};
    sprintf(input_str, "%s", env->GetStringUTFChars(input_url, JNI_FALSE));
    sprintf(output_str, "%s", env->GetStringUTFChars(output_url, JNI_FALSE));

    //FFmpeg av_log() callback
    av_log_set_callback(custom_log);

    av_register_all();
    //Network
    avformat_network_init();

    int videoindex = -1;
    int64_t start_time;

    //Input
    char buf[1024];
    if ((ret = avformat_open_input(&ifmt_ctx, input_str, 0, 0)) < 0) {
        av_strerror(ret, buf, 1024);
        jstring msg = env->NewStringUTF("Could not open input file.");
        env->CallVoidMethod(resultObj, jmethodId_Result, ret, msg);
        LOGE("Could not open input file.");
        goto end;
    }
    if ((ret = avformat_find_stream_info(ifmt_ctx, 0)) < 0) {
        jstring msg = env->NewStringUTF("Failed to retrieve input stream information");
        env->CallVoidMethod(resultObj, jmethodId_Result, ret, msg);
        LOGE("Failed to retrieve input stream information");
        goto end;
    }


    for (i = 0; i < ifmt_ctx->nb_streams; i++) {
        if (ifmt_ctx->streams[i]->codec->codec_type == AVMEDIA_TYPE_VIDEO) {
            videoindex = i;
            break;
        }
    }
    //Output
    avformat_alloc_output_context2(&ofmt_ctx, nullptr, "flv", output_str); //RTMP
    //avformat_alloc_output_context2(&ofmt_ctx, nullptr, "mpegts", output_str);//UDP

    if (!ofmt_ctx) {
        ret = AVERROR_UNKNOWN;
        jstring msg = env->NewStringUTF("Could not create output context");
        env->CallVoidMethod(resultObj, jmethodId_Result, ret, msg);
        LOGE("Could not create output context\n");
        goto end;
    }
    ofmt = ofmt_ctx->oformat;
    for (i = 0; i < ifmt_ctx->nb_streams; i++) {
        //Create output AVStream according to input AVStream
        AVStream *in_stream = ifmt_ctx->streams[i];
        AVStream *out_stream = avformat_new_stream(ofmt_ctx, in_stream->codec->codec);
        if (!out_stream) {
            ret = AVERROR_UNKNOWN;
            jstring msg = env->NewStringUTF("Failed allocating output stream");
            env->CallVoidMethod(resultObj, jmethodId_Result, ret, msg);
            LOGE("Failed allocating output stream\n");

            goto end;
        }
        //Copy the settings of AVCodecContext
        ret = avcodec_copy_context(out_stream->codec, in_stream->codec);
        if (ret < 0) {
            jstring msg = env->NewStringUTF(
                    "Failed to copy context from input to output stream codec context");
            env->CallVoidMethod(resultObj, jmethodId_Result, ret, msg);
            LOGE("Failed to copy context from input to output stream codec context\n");
            goto end;
        }
        out_stream->codec->codec_tag = 0;
        if (ofmt_ctx->oformat->flags & AVFMT_GLOBALHEADER)
            out_stream->codec->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;

    }

    //Open output URL
    if (!(ofmt->flags & AVFMT_NOFILE)) {
        ret = avio_open(&ofmt_ctx->pb, output_str, AVIO_FLAG_WRITE);
        if (ret < 0) {
            char *p = "Could not open output URL:";
            strcat(p, output_str);
            jstring msg = env->NewStringUTF(p);
            env->CallVoidMethod(resultObj, jmethodId_Result, ret, msg);
            LOGE("Could not open output URL '%s'", output_str);
            goto end;
        }
    }
    //Write file header
    ret = avformat_write_header(ofmt_ctx, nullptr);
    if (ret < 0) {
        jstring msg = env->NewStringUTF("Error occurred when opening output URL");
        env->CallVoidMethod(resultObj, jmethodId_Result, ret, msg);
        LOGE("Error occurred when opening output URL\n");
        goto end;
    }
    LOGE("Success!\n");
    start_time = av_gettime();

    while (true) {
        int frame_index = 0;
        AVStream *in_stream, *out_stream;
        //Get an AVPacket
        ret = av_read_frame(ifmt_ctx, &pkt);
        if (ret < 0) {
            jstring msg = env->NewStringUTF("av_read_frame failed");
            env->CallVoidMethod(resultObj, jmethodId_Result, ret, msg);
            break;
        }
        //FIX：No PTS (Example: Raw H.264)
        //Simple Write PTS
        if (pkt.pts == AV_NOPTS_VALUE) {
            //Write PTS
            AVRational time_base1 = ifmt_ctx->streams[videoindex]->time_base;
            //Duration between 2 frames (us)
            int64_t calc_duration =
                    (double) AV_TIME_BASE / av_q2d(ifmt_ctx->streams[videoindex]->r_frame_rate);
            //Parameters
            pkt.pts = (double) (frame_index * calc_duration) /
                      (double) (av_q2d(time_base1) * AV_TIME_BASE);
            pkt.dts = pkt.pts;
            pkt.duration =
                    (double) calc_duration / (double) (av_q2d(time_base1) * AV_TIME_BASE);
        }
        //Important:Delay
        if (pkt.stream_index == videoindex) {
            AVRational time_base = ifmt_ctx->streams[videoindex]->time_base;
            AVRational time_base_q = {1, AV_TIME_BASE};
            int64_t pts_time = av_rescale_q(pkt.dts, time_base, time_base_q);
            int64_t now_time = av_gettime() - start_time;
            if (pts_time > now_time) {
                av_usleep(pts_time - now_time);
            }
        }

        in_stream = ifmt_ctx->streams[pkt.stream_index];
        out_stream = ofmt_ctx->streams[pkt.stream_index];
        /* copy packet */
        //Convert PTS/DTS
        pkt.pts = av_rescale_q_rnd(pkt.pts, in_stream->time_base, out_stream->time_base,
                                   AV_ROUND_NEAR_INF);
        pkt.dts = av_rescale_q_rnd(pkt.dts, in_stream->time_base, out_stream->time_base,
                                   AV_ROUND_NEAR_INF);
        pkt.duration = av_rescale_q(pkt.duration, in_stream->time_base, out_stream->time_base);
        pkt.pos = -1;
        //Print to Screen
        if (pkt.stream_index == videoindex) {
            LOGE("Send %8d video frames to output URL\n", frame_index);
            frame_index++;
        }
        //ret = av_write_frame(ofmt_ctx, &pkt);
        ret = av_interleaved_write_frame(ofmt_ctx, &pkt);

        if (ret < 0) {
            jstring msg = env->NewStringUTF("Error muxing packet");
            env->CallVoidMethod(resultObj, jmethodId_Result, ret, msg);
            LOGE("Error muxing packet\n");
            break;
        }
        av_free_packet(&pkt);
    }
    {
        jstring msg = env->NewStringUTF("publish end！");
        env->CallVoidMethod(resultObj, jmethodId_Result, ret, msg);
    }
    //Write file trailer
    av_write_trailer(ofmt_ctx);
    end:
    avformat_close_input(&ifmt_ctx);
    /* close output */
    if (ofmt_ctx && !(ofmt->flags & AVFMT_NOFILE))
        avio_close(ofmt_ctx->pb);
    avformat_free_context(ofmt_ctx);
    if (ret < 0 && ret != AVERROR_EOF) {
        jstring msg = env->NewStringUTF("Error occurred.");
        env->CallVoidMethod(resultObj, jmethodId_Result, ret, msg);
        LOGE("Error occurred.\n");
        return -1;
    }
    return 0;
}